#include "image_collector.h"
#include "camera_receiver.h"

#include "muduo/base/CountDownLatch.h"
#include "muduo/base/Timestamp.h"

#include <iomanip>

#include <cv_bridge/cv_bridge.h>

using namespace sensor::camera;

namespace {

std::string formatTimestamp(const muduo::Timestamp& ts)
{
    char buf[64] = { 0 };
    time_t seconds = ts.microSecondsSinceEpoch() / 1000 / 1000;
    struct tm tm_time;
    gmtime_r(&seconds, &tm_time);

    snprintf(buf, sizeof(buf), "%4d%02d%02d_%02d%02d%02d",
        tm_time.tm_year + 1900, tm_time.tm_mon + 1, tm_time.tm_mday, tm_time.tm_hour,
        tm_time.tm_min, tm_time.tm_sec);
    return buf;
}

}

ImageCollector::ImageCollector(size_t devnum)
    : rclcpp::Node("ImageCollector")
    , devnum_(devnum)
    , logger_(rclcpp::get_logger("ImageCollector"))
{
    pub_ = this->create_publisher<interface::msg::Images>("/sensor/camera/images", 30);
}

void ImageCollector::infuseOneFrame(const std::string& device, const std::shared_ptr<CameraOneFrame> frame)
{
    std::unordered_map<std::string, std::shared_ptr<CameraOneFrame>> output;
    {
        std::lock_guard<std::mutex> guard(frames_lock_);
        fps_[device].receivedOneFrame();
        frames_[device] = frame;
        if (devnum_ != frames_.size()) {
            return;
        }
        output.swap(frames_);
    }

    auto now = muduo::Timestamp::now();

    if (!threads_) {
        threads_ = std::make_unique<muduo::ThreadPool>("image_transfer");
        threads_->start(8);
    }

    muduo::CountDownLatch cdl(output.size());
    std::set<sensor_msgs::msg::Image::SharedPtr> images;
    for (auto iter = output.begin(); iter != output.end(); iter++) {
        const auto key = iter->first;
        const auto image = iter->second->image;
        auto worker = std::bind(&ImageCollector::processOneFrame, this, now.microSecondsSinceEpoch(), key, image, std::ref(images), std::ref(cdl));

        // 保存2K图片单路耗时就很大, 使用线程池优化不明显，需要优化memcpy本身的性能，aarch64支持SIMD，可以提升一些性能，待尝试
        // 对于保存视频，可以尝试 线程池优化
        threads_->run(worker);
        // worker();
    }

    cdl.wait();

    interface::msg::Images images_tobe_send;
    images_tobe_send.header.stamp.sec = now.secondsSinceEpoch();
    images_tobe_send.header.stamp.nanosec = (now.microSecondsSinceEpoch() % 1000000) * 1000;
    // images_tobe_send.index = image_tx_counter_++;
    for (auto& image : images) {
        images_tobe_send.images.emplace_back(std::move(*image));
    }

    pub_->publish(images_tobe_send);

    // 利用VideoWriter保存mp4, 每5分钟一个片段.
    if (muduo::timeDifference(now, last_video_saved_time_) > 5 * 60) {
        std::lock_guard<std::mutex> guard(video_writer_lock_);
        video_writers_.clear();
        last_video_saved_time_ = now;
    }

    std::stringstream ss;
    for (auto& fps : fps_) {
        ss << std::fixed << std::setprecision(2) << fps.first << ":" << fps.second.fps() << " ";
    }

    std::stringstream aligned_ss;
    for (auto& fps : aligned_fps_) {
        aligned_ss << std::fixed << std::setprecision(2) << fps.first << ":" << fps.second.fps() << " ";
    }

    //    RCLCPP_INFO(logger_, "Collect frames fps:%s", ss.str().c_str());
    RCLCPP_INFO(logger_, "Collect all camera frames and published!, fps:%s", aligned_ss.str().c_str());
}

void ImageCollector::reset()
{
    // trigger_index_coutner_++;
    // RCLCPP_INFO(logger_, "ImageCollector start! %d", trigger_index_coutner_);
    // std::lock_guard<std::mutex> guard(lock_);
    // frames_.clear();
}

void ImageCollector::processOneFrame(int64_t now, const std::string& key, cv::Mat& in, std::set<sensor_msgs::msg::Image::SharedPtr>& images,
    muduo::CountDownLatch& cdl)
{
    //    RCLCPP_INFO(logger_, "key %s:%d start make ros message", key.c_str(), publish_index_counter_);

    auto header = std_msgs::msg::Header();
    header.stamp.sec = now / 1000000;
    header.stamp.nanosec = (now % 1000000) * 1000;

    header.frame_id = key;

    auto msg = cv_bridge::CvImage(header, "bgr8", in).toImageMsg();
    {
        std::lock_guard<std::mutex> guard(frames_lock_);
        images.insert(msg);
        aligned_fps_[key].receivedOneFrame(); // for statistics
    }

    //    RCLCPP_INFO(logger_, "key %s:%d make ros message done and start save video", key.c_str(), publish_index_counter_);
#if 0
    // 保存mp4
    std::shared_ptr<cv::VideoWriter> writer;
    {
        std::lock_guard<std::mutex> guard(video_writer_lock_);
        auto iter = video_writers_.find(key);
        if (iter == video_writers_.end()) {
            const std::string prefix = "/home/nvidia/workspace/data/videos/";
            const auto fname = prefix + formatTimestamp(muduo::Timestamp::now()) + "_" + key + ".mp4";
            video_writers_[key] = std::make_shared<cv::VideoWriter>(fname, cv::VideoWriter::fourcc('m', 'p', '4', 'v'), 30, in.size());
        }
        writer = video_writers_[key];
    }

    // 程序运行中ctrl-c或者其原因意外推出，VideoWriter销毁不规范，会导致正在保存的视频不可用. 后续可优化
    writer->write(in);
#endif
    // RCLCPP_INFO(logger_, "key %s:%d save video done", key.c_str(), publish_index_counter_);
#if 0
    auto simage = in.clone();
    cv::resize(simage, simage, { 1920, 1080 });
    cv::imshow(key, simage);
    cv::waitKey(10);
#endif

    cdl.countDown();
}
